# 项目: hrv_predict_project
# 目录结构:
# hrv_predict_project/
# ├── data/
# │   └── patients_hrv.json
# ├── result/
# │   ├── report/
# │   └── pic/
# └── src/
#     ├── data_preprocess.py
#     ├── model_logistic.py
#     ├── model_mlp.py
#     ├── model_rf.py
#     ├── model_lgbm.py
#     ├── model_xgb.py
#     ├── utils.py
#     └── train_eval.py

# ---------------------------
# File: src/data_preprocess.py
# ---------------------------
import os
import json
import pandas as pd
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.model_selection import train_test_split

def load_and_preprocess(json_path, test_size=0.2, random_state=0):
    # 读取 JSON 数据
    df = pd.read_json(json_path)
    # 指定 HRV 特征列（根据实际调整）
    hrv_features = [c for c in df.columns if c not in ['age', 'Gender', 'AnalysisText']]

    # 缺失值填充（均值）
    imputer = SimpleImputer(strategy='mean')
    df[hrv_features] = imputer.fit_transform(df[hrv_features])

    # 特征标准化
    scaler = StandardScaler()
    df[hrv_features] = scaler.fit_transform(df[hrv_features])

    # 构建标签
    df['age_group'] = ((df['age'] // 10) * 10).astype(int)
    le = LabelEncoder()
    df['gender_label'] = le.fit_transform(df['Gender'])

    # 特征矩阵和标签
    X = df[hrv_features].values
    y_age = df['age_group'].values
    y_gender = df['gender_label'].values

    # 数据拆分
    X_tr_age, X_te_age, y_tr_age, y_te_age = train_test_split(
        X, y_age, test_size=test_size, random_state=random_state, stratify=y_age)
    X_tr_gen, X_te_gen, y_tr_gen, y_te_gen = train_test_split(
        X, y_gender, test_size=test_size, random_state=random_state, stratify=y_gender)

    return (X_tr_age, X_te_age, y_tr_age, y_te_age,
            X_tr_gen, X_te_gen, y_tr_gen, y_te_gen,
            hrv_features)

# ---------------------------
# File: src/model_logistic.py
# ---------------------------
import torch
import torch.nn as nn

class LogisticRegressionModel(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.linear = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        logits = self.linear(x)
        if logits.shape[1] == 1:
            return torch.sigmoid(logits)
        else:
            return torch.softmax(logits, dim=1)

# ---------------------------
# File: src/model_mlp.py
# ---------------------------
import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, output_dim)
        )

    def forward(self, x):
        out = self.net(x)
        if out.shape[1] == 1:
            return torch.sigmoid(out)
        else:
            return torch.softmax(out, dim=1)

# ---------------------------
# File: src/model_rf.py
# ---------------------------
from sklearn.ensemble import RandomForestClassifier

def get_rf_model(n_estimators=100, random_state=0):
    return RandomForestClassifier(n_estimators=n_estimators, random_state=random_state)

# ---------------------------
# File: src/model_lgbm.py
# ---------------------------
import lightgbm as lgb

def get_lgbm_model(n_estimators=100, random_state=0):
    return lgb.LGBMClassifier(n_estimators=n_estimators,
                              random_state=random_state, importance_type='gain')

# ---------------------------
# File: src/model_xgb.py
# ---------------------------
import xgboost as xgb

def get_xgb_model(n_estimators=100, random_state=0):
    return xgb.XGBClassifier(n_estimators=n_estimators,
                             random_state=random_state,
                             use_label_encoder=False,
                             eval_metric='logloss')

# ---------------------------
# File: src/utils.py
# ---------------------------
import os
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np


def save_classification_report(report_dict, filepath):
    df = pd.DataFrame(report_dict).transpose()
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    df.to_csv(filepath, index=True)


def plot_feature_importance(importances, feature_names, out_path, top_n=10, title=None):
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    idx = np.argsort(importances)[::-1][:top_n]
    top_feats = [feature_names[i] for i in idx]
    top_imp = importances[idx]
    plt.figure(figsize=(8, 6))
    plt.barh(range(top_n), top_imp[::-1])
    plt.yticks(range(top_n), top_feats[::-1])
    plt.xlabel('Importance')
    if title:
        plt.title(title)
    plt.tight_layout()
    plt.savefig(out_path)
    plt.close()

# ---------------------------
# File: src/train_eval.py
# ---------------------------
import os
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.metrics import classification_report
import pandas as pd
from data_preprocess import load_and_preprocess
from model_logistic import LogisticRegressionModel
from model_mlp import MLP
from model_rf import get_rf_model
from model_lgbm import get_lgbm_model
from model_xgb import get_xgb_model
from utils import save_classification_report, plot_feature_importance


def train_pytorch_model(model, X_train, y_train, X_test, y_test, task_name, out_report):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    # 损失和优化
    if model.linear.out_features == 1:
        criterion = nn.BCELoss()
    else:
        criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=1e-3)

    # 转为 Tensor
    X_tr = torch.FloatTensor(X_train).to(device)
    X_te = torch.FloatTensor(X_test).to(device)
    y_tr = torch.LongTensor(y_train).to(device)

    # 训练
    model.train()
    for epoch in range(50):
        optimizer.zero_grad()
        outputs = model(X_tr)
        if outputs.shape[1] == 1:
            loss = criterion(outputs.view(-1), y_tr.float())
        else:
            loss = criterion(outputs, y_tr)
        loss.backward()
        optimizer.step()

    # 预测
    model.eval()
    with torch.no_grad():
        preds = model(X_te)
        if preds.shape[1] == 1:
            y_pred = (preds.view(-1) > 0.5).long().cpu().numpy()
        else:
            y_pred = torch.argmax(preds, dim=1).cpu().numpy()

    # 报告
    report = classification_report(y_test, y_pred, output_dict=True)
    save_classification_report(report, out_report)


def train_sklearn_model(ModelClass, X_train, y_train, X_test, y_test,
                        task_name, out_report, feat_names, pic_path=None):
    model = ModelClass()
    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)
    report = classification_report(y_test, y_pred, output_dict=True)
    save_classification_report(report, out_report)
    # 特征重要性
    if hasattr(model, 'feature_importances_') and pic_path:
        plot_feature_importance(
            model.feature_importances_, feat_names,
            pic_path, title=f"{task_name} Feature Importance")


def main():
    # 数据预处理
    json_path = os.path.join('..', 'data', 'patients_hrv.json')
    (X_tr_age, X_te_age, y_tr_age, y_te_age,
     X_tr_gen, X_te_gen, y_tr_gen, y_te_gen,
     feature_names) = load_and_preprocess(json_path)

    # 输出路径
    report_dir = os.path.join('..', 'result', 'report')
    pic_dir = os.path.join('..', 'result', 'pic')

    # 任务：年龄段预测
    # PyTorch Logistic
    lr_age = LogisticRegressionModel(input_dim=X_tr_age.shape[1], output_dim=len(set(y_tr_age)))
    train_pytorch_model(lr_age, X_tr_age, y_tr_age, X_te_age, y_te_age,
                        'age_logistic', os.path.join(report_dir, 'age_logistic.csv'))
    # PyTorch MLP
    mlp_age = MLP(input_dim=X_tr_age.shape[1], hidden_dim=64, output_dim=len(set(y_tr_age)))
    train_pytorch_model(mlp_age, X_tr_age, y_tr_age, X_te_age, y_te_age,
                        'age_mlp', os.path.join(report_dir, 'age_mlp.csv'))
    # Random Forest
    train_sklearn_model(get_rf_model, X_tr_age, y_tr_age, X_te_age, y_te_age,
                        'age_rf', os.path.join(report_dir, 'age_rf.csv'),
                        feature_names, os.path.join(pic_dir, 'age_rf.png'))
    # LightGBM
    train_sklearn_model(get_lgbm_model, X_tr_age, y_tr_age, X_te_age, y_te_age,
                        'age_lgbm', os.path.join(report_dir, 'age_lgbm.csv'),
                        feature_names, os.path.join(pic_dir, 'age_lgbm.png'))
    # XGBoost
    train_sklearn_model(get_xgb_model, X_tr_age, y_tr_age, X_te_age, y_te_age,
                        'age_xgb', os.path.join(report_dir, 'age_xgb.csv'),
                        feature_names, os.path.join(pic_dir, 'age_xgb.png'))

    # 任务：性别预测
    # PyTorch Logistic
    lr_gen = LogisticRegressionModel(input_dim=X_tr_gen.shape[1], output_dim=1)
    train_pytorch_model(lr_gen, X_tr_gen, y_tr_gen, X_te_gen, y_te_gen,
                        'gender_logistic', os.path.join(report_dir, 'gender_logistic.csv'))
    # PyTorch MLP
    mlp_gen = MLP(input_dim=X_tr_gen.shape[1], hidden_dim=64, output_dim=1)
    train_pytorch_model(mlp_gen, X_tr_gen, y_tr_gen, X_te_gen, y_te_gen,
                        'gender_mlp', os.path.join(report_dir, 'gender_mlp.csv'))
    # Random Forest
    train_sklearn_model(get_rf_model, X_tr_gen, y_tr_gen, X_te_gen, y_te_gen,
                        'gender_rf', os.path.join(report_dir, 'gender_rf.csv'),
                        feature_names, os.path.join(pic_dir, 'gender_rf.png'))
    # LightGBM
    train_sklearn_model(get_lgbm_model, X_tr_gen, y_tr_gen, X_te_gen, y_te_gen,
                        'gender_lgbm', os.path.join(report_dir, 'gender_lgbm.csv'),
                        feature_names, os.path.join(pic_dir, 'gender_lgbm.png'))
    # XGBoost
    train_sklearn_model(get_xgb_model, X_tr_gen, y_tr_gen, X_te_gen, y_te_gen,
                        'gender_xgb', os.path.join(report_dir, 'gender_xgb.csv'),
                        feature_names, os.path.join(pic_dir, 'gender_xgb.png'))

if __name__ == '__main__':
    main()
